from typing import List, Tuple
import numpy as np


class GameRules:
    """游戏规则管理器 - 处理五子棋基本规则"""
    
    def __init__(self, board_size: int):
        self.board_size = board_size
    
    def is_valid_move(self, board: np.ndarray, action: int) -> bool:
        """检查动作是否合法"""
        row, col = self._action_to_position(action)
        return (0 <= row < self.board_size and 
                0 <= col < self.board_size and 
                board[row, col] == 0)
    
    def get_legal_actions(self, board: np.ndarray) -> List[int]:
        """获取所有合法动作"""
        legal_actions = []
        for i in range(self.board_size):
            for j in range(self.board_size):
                if board[i, j] == 0:
                    legal_actions.append(i * self.board_size + j)
        return legal_actions
    
    def check_win(self, board: np.ndarray, row: int, col: int, player: int) -> bool:
        """检查是否获胜"""
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        
        for dr, dc in directions:
            if self._count_consecutive(board, row, col, dr, dc, player) >= 5:
                return True
        return False
    
    def is_board_full(self, board: np.ndarray) -> bool:
        """检查棋盘是否已满"""
        return not np.any(board == 0)
    
    def _action_to_position(self, action: int) -> Tuple[int, int]:
        """将动作转换为位置坐标"""
        return action // self.board_size, action % self.board_size
    
    def _count_consecutive(self, board: np.ndarray, row: int, col: int, 
                          dr: int, dc: int, player: int) -> int:
        """计算连续棋子数量"""
        count = 1  # 包含当前落子
        
        # 正方向
        r, c = row + dr, col + dc
        while (0 <= r < self.board_size and 0 <= c < self.board_size and 
               board[r, c] == player):
            count += 1
            r, c = r + dr, c + dc
        
        # 反方向
        r, c = row - dr, col - dc
        while (0 <= r < self.board_size and 0 <= c < self.board_size and 
               board[r, c] == player):
            count += 1
            r, c = r - dr, c - dc
        
        return count